import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.ensemble import ExtraTreesClassifier
from sklearn.metrics import accuracy_score, precision_score, recall_score
from gen_corr import get_corr_feature


train_df = pd.read_csv('data_c1.csv')
test_df_4 = pd.read_csv('data_c4.csv')
test_df_6 = pd.read_csv('data_c6.csv')
cols = get_corr_feature()

X_train = train_df.drop('label', axis=1)[cols].values
y_train = train_df['label'].values

X_test_4 = test_df_4.drop('label', axis=1)[cols].values
y_test_4 = test_df_4['label'].values

X_test_6 = test_df_6.drop('label', axis=1)[cols].values
y_test_6 = test_df_6['label'].values

clf = ExtraTreesClassifier(random_state=42, n_estimators=100)
clf.fit(X_train, y_train)

y_pred_4 = clf.predict(X_test_4)

accuracy_4 = accuracy_score(y_test_4, y_pred_4)

print(f'Accuracy_4: {accuracy_4:.4f}')

y_pred_6 = clf.predict(X_test_6)
accuracy_6 = accuracy_score(y_test_6, y_pred_6)
print(f'Accuracy_6: {accuracy_6:.4f}')
